import java.util.HashMap;
import java.util.Map;
import java.util.Iterator;
import java.util.Arrays;

/**
 * ContextDataCollection class tabulates number of zero and one bits corresponding to
 * each context characteristics.
 *  modified by edhong to include splitting the data collection. 
 * 
 * @author Yong Shin
 *
 */
public class ContextDataCollection {

	private HashMap<ContextCharacteristics, BitCounts> bitCountMap;
        private BitCounts totalbits;
    private CCBitCounts[] sortedccbit; // sortedlist of CCBitCounts;
    // only updated when computeOptimalAdaptiveCost() called. 
    private int[] optsplits; // dyn prog solution array used in computeOptimalAdaptiveCost. 
    private double[] optcosts; // dyn prog solution array used in computeOptimalAdaptiveCost. 
	
	public ContextDataCollection() {
		bitCountMap = new HashMap<ContextCharacteristics, BitCounts>();
                totalbits = new BitCounts(0,0);
		sortedccbit = null;
		optsplits = null;
		optcosts= null;
	}

        // copy constructor that does not create new underlying Bitcounts.
        public ContextDataCollection(ContextDataCollection other) {
		bitCountMap = new HashMap<ContextCharacteristics, BitCounts>(other.bitCountMap);
                totalbits = new BitCounts(other.totalbits);
		sortedccbit = other.sortedccbit;
		optsplits = other.optsplits;
		optcosts = other.optcosts;
        }
        
        public BitCounts get(ContextCharacteristics c) {
            return bitCountMap.get(c);
        }


	/**
	 * Adds bit counts for specified context characteristics.
	 * @param cc
	 * @param zeroBits
	 * @param oneBits
	 */
	public void addBitCount(ContextCharacteristics cc, int zeroBits, int oneBits) {
                totalbits.oneBits += oneBits;
                totalbits.zeroBits += zeroBits;
		if (bitCountMap.containsKey(cc)) {
			BitCounts bcounts = bitCountMap.get(cc);
			bcounts.oneBits += oneBits;
			bcounts.zeroBits += zeroBits;
		} else {
			BitCounts bcounts = new BitCounts(zeroBits, oneBits);
			bitCountMap.put(cc, bcounts);
		}
	}



	/**
	 * Returns bit counts corresponding to a context characteristics.
	 * @return 
	 */
	public BitCounts getTotal() {
		return new BitCounts(totalbits);
	}
	
	/**
	 * Returns total bit counts for given context id. Bit counts for context values whose bucket
	 * id's are <= given bucket id is stored in BitCount[0].  Bit counts for context values whose
	 * bucket id's are > given bucket id is stored in BitCount[1].
	 * @param contextID
	 * @return
	 */
	public BitCounts[] getBitCounts(int contextID, int bucketID) {
		BitCounts[] result = new BitCounts[2];
		result[0] = new BitCounts(0,0);
		result[1] = new BitCounts(0,0);
		Iterator<ContextCharacteristics> ccs = bitCountMap.keySet().iterator();
		while (ccs.hasNext()) {
			ContextCharacteristics cc = ccs.next();
			BitCounts bcounts = bitCountMap.get(cc);
			if (cc.getBucketID(contextID) <= bucketID) {
				result[0].add(bcounts);
			} else {
				result[1].add(bcounts);
			}
		}
		return result;
	}

	/**
	 * Calculates the percentage of the part with respect to the whole.
	 * @param part
	 * @param whole
	 * @return
	 */
	private int percent(int part, int whole) {
		return (int) Math.round(part * 100.0 / whole);
	}

        public void printBitDistributionStats() {
        			int numCtx = ContextDefinition.getContextDefinition().getNumContexts();

				System.out.println("Bit value distributions (zerobit onebit):");
				for (int i=0; i<numCtx; i++) {
					BitCounts prev = new BitCounts(0,0);
					Context ctx = ContextDefinition.getContextDefinition().getContext(i);
					System.out.println("Context " + ctx.toString());
					int numBuc = ctx.getNumberBuckets();
					for (int j=0; j<numBuc; j++) {
						BitCounts[] bcs = this.getBitCounts(i, j);
						int zb = bcs[0].zeroBits - prev.zeroBits;
						int ob = bcs[0].oneBits - prev.oneBits;
						int tb = zb + ob;
						System.out.println("\t(" + (j+1) + ") " + 
								zb + " " + ob + " [" + percent(zb,tb) + "% " + percent(ob,tb) + "%]");
						prev = bcs[0];
					}
				}
        }

	/**
	 * Removes some of the data stored in this ContextDataCollection. These data are 
         * put into a newly created ContextDataCollection. Returns the new ContextDataCollection. 
         * The items that remain in this collection are those whose bucket ID for the given context
         * ID are <= the given bucket ID value. Items moved over are those whose
         * bucket ID for the given context ID are > the given bucket ID value. 
	 * @param contextID
         * @param bucketID
	 * @return
	 */
	public ContextDataCollection split(int contextID, int bucketID) {
                ContextDataCollection newdata = new ContextDataCollection();
		Iterator<ContextCharacteristics> ccs = bitCountMap.keySet().iterator();
		while (ccs.hasNext()) {
			ContextCharacteristics cc = ccs.next();
			if (cc.getBucketID(contextID) <= bucketID) {
                            // data remains, do nothing. 
                        } else {
                            BitCounts bcounts = bitCountMap.get(cc);                          
                            newdata.addBitCount(cc,bcounts.zeroBits, bcounts.oneBits);
                            totalbits.subtract(bcounts);
                            ccs.remove();
			}
		}
		return newdata;
	}

	/**
	 * Returns the number of data in the collection.
	 * @return
	 */
	public int size() {
		return bitCountMap.size();
	}
	
	/**
	 * Returns the iterator for the context characteristics.
	 * @return
	 */
	public Iterator<ContextCharacteristics> ccIterator() {
		return bitCountMap.keySet().iterator();
	}
 
        public String toString() {
            return "totalbits: " + totalbits  + "\n" + bitCountMap + "\n";
        }

    /** 
     *  computes and returns first-order entropy of underlying bitcounts, each considered separately
     */
    public double getEntropy() {
	double retval = 0.0F;
	for (BitCounts bc : bitCountMap.values()) {
	    retval += bc.getEntropy();
	}
	return retval;
    }

    /** 
     *  computes and returns adaptive Encoding Cost of underlying bitcounts, each considered separately
     */
    public double getAdaptiveCost() {
	double retval = 0.0F;
	for (BitCounts bc : bitCountMap.values()) {
	    retval += bc.getAdaptiveCost();
	}
	return retval;	
    }

    /** 
     *  prints out optimal solution grouping of bitcounts (no tree restriction)
     *    to minimize adaptive encoding cost. 
     *  calls computeOptimalAdpativeCost(). 
     */
    
    public void computeAndPrintOptimalAdaptiveCostSolution() {
	computeOptimalAdaptiveCost();
	System.out.println("Optimal nontree cost is " + 
			   optcosts[size()-1]);
	int groupcount = 0;
	int i=size()-1;
	while (i>0) { 
	    double subsolutioncost;
	    if (optsplits[i]==-1)
		subsolutioncost = 0.0f;
	    else 
		subsolutioncost = optcosts[optsplits[i]];
	    System.out.println("Group " + groupcount + " cost: " + 
			       (optcosts[i] - subsolutioncost));
	    for (int j=i; j > optsplits[i];  j--) {
		System.out.println( " "+ groupcount + ":" + sortedccbit[j]);
	    }
	    i = optsplits[i];
	    groupcount++;
	}
    }

 
    // contracts data from sortedccbit into 
    //  smallsortedccbit and returns # valid items in smallsortedccbit. 
    //  contextchar fields of smallsortedccbit objects set to null.
    // smalltofull_index[i] is set to contain index of last item in sortedccbit
    // that belongs to the collapsed group smallsortedccbit[i];
    private int collapse(CCBitCounts[] smallsortedccbit, int[] smalltofull_index) {
	BitCounts current;
	current = new BitCounts(sortedccbit[0].bitCount);
	smallsortedccbit[0] = new CCBitCounts(null, current);
	int smallindex = 0;
	for (int i=1; i < size(); i++) {
	    if (current.ratioEquals(sortedccbit[i].bitCount)) {
		current.add(sortedccbit[i].bitCount);
	    }
	    else {
		smalltofull_index[smallindex] = i-1;
		current = new BitCounts(sortedccbit[i].bitCount);		
		smallindex++;
		smallsortedccbit[smallindex] = new CCBitCounts(null, current);
	    }
	}
	smalltofull_index[smallindex] = size()-1; // set index for last group. 
	
	return smallindex+1;
    }

    // expands data from  smalloptsplits, and smalloptcosts  into 
    //  optsplits, and optcosts. does not fill in not-computed values. 
    private void expand(int[] smalltofull_index, int smallsize, 
			int[] smalloptsplits, double[] smalloptcosts) {

	for (int i=0; i < smallsize; i++) {
	    int j = smalltofull_index[i];
	    if (smalloptsplits[i] == -1)
		optsplits[j] = -1;
	    else 
		optsplits[j] = smalltofull_index[smalloptsplits[i]];
	    optcosts[j] = smalloptcosts[i];
	}

    }

    

    /** 
     *  computes, stores, and returns optimal grouping of bitcounts (no tree restriction)
     *    to minimize adaptive encoding cost. 
     */
    public double computeOptimalAdaptiveCost() {
	sortedccbit = new CCBitCounts[size()];
	Iterator<Map.Entry<ContextCharacteristics,BitCounts>> it; 
	it = bitCountMap.entrySet().iterator();
	for (int i=0; i < size(); i++) {
	    Map.Entry<ContextCharacteristics,BitCounts> x;
	    x = it.next();
	    sortedccbit[i] = new CCBitCounts(x.getKey(), x.getValue()); 
	}

	Arrays.sort(sortedccbit);

	CCBitCounts[] smallsortedccbit = new CCBitCounts[size()];
	int[] smalltofull_index = new int[size()];
	int smallsize =0;
	
	smallsize = collapse(smallsortedccbit, smalltofull_index);

	/*
	//debugging print statements.
		BitCounts tempbp;	
		System.out.println("sortedccbit:");
	tempbp = new BitCounts(0,0);
	for (int i=size()-1; i >=0; i--) {
	    System.out.print(sortedccbit[i]);
	    tempbp.add(sortedccbit[i].bitCount);
	    System.out.println(" " + tempbp);
	    }


	System.out.println("smalltofull_index:" + Arrays.toString(smalltofull_index));

	System.out.println("smallsortedccbit:");
	tempbp = new BitCounts(0,0);
	for (int i=smallsize-1; i >=0; i--) {
	    System.out.print(smallsortedccbit[i]);
	    tempbp.add(smallsortedccbit[i].bitCount);
	    System.out.println(" " + tempbp);
	    }

	*/

	// optsplits[k] contains index of where last group ended in optimal split for [1...k].
	//   -1 == no previous group. 
	optsplits = new int[size()];
	// optcosts[k] contains optimal cost for group bitcounts 1...k
	optcosts = new double[size()];

	int[] smalloptsplits = new int[smallsize];
	double[] smalloptcosts = new double[smallsize];
	// initialize
	    smalloptsplits[0] = -1;
	    smalloptcosts[0] = smallsortedccbit[0].bitCount.getAdaptiveCost();

	for (int i=1; i < smallsize; i++) {
	    smalloptsplits[i] = -1;
	    smalloptcosts[i] = smalloptcosts[i-1] + smallsortedccbit[i].bitCount.getAdaptiveCost();
	}
	
	for (int i=0; i < smallsize; i++) {	    
	    //	    if (i%10==0) System.out.print(".");
	    BitCounts bc = new BitCounts(0,0);
	    for (int j=i; j >=1 ; j--) {
		bc.add(smallsortedccbit[j].bitCount);
		double newsol = bc.getAdaptiveCost() + smalloptcosts[j-1];
		if (newsol < smalloptcosts[i]) {
		    smalloptcosts[i] = newsol;
		    smalloptsplits[i] = j-1;
		}
	    }
	    bc.add(smallsortedccbit[0].bitCount);
	    if (bc.getAdaptiveCost() < smalloptcosts[i]) {
		smalloptcosts[i] = bc.getAdaptiveCost();
		smalloptsplits[i] = -1;
	    }
	}

	/*
	// debugging print statements...
	System.out.print("Costs(small):");
	System.out.println(Arrays.toString(smalloptcosts));
	System.out.print("Splits(small):");
	System.out.println(Arrays.toString(smalloptsplits));
	*/

	expand( smalltofull_index, smallsize, smalloptsplits, smalloptcosts);

	/*
	// debugging print statements...
	System.out.print("Costs:");
	System.out.println(Arrays.toString(optcosts));
	System.out.print("Splits:");
	System.out.println(Arrays.toString(optsplits));
	*/

	return optcosts[size()-1];
    }

    // gets rid of dyn prog solutions. 
    public void freedata() {
	optcosts = null;
	optsplits = null;
	sortedccbit= null;
    }

}

 class CCBitCounts implements Comparable<CCBitCounts> {
    ContextCharacteristics contextChar;
    BitCounts bitCount;
    public CCBitCounts(ContextCharacteristics cc, BitCounts bc) {
	contextChar = cc;
	bitCount = bc;
    }
    
    @Override
	public boolean equals(Object other) {
	if (other instanceof CCBitCounts) {
	    CCBitCounts o = (CCBitCounts) other;
	    return (this.contextChar.equals(o.contextChar)
		    &&
		    this.bitCount.equals(o.bitCount));
	}
	else return false;
    }

    // compares primarily by  bitcount, secondaril by contextCharacteristic.
    public int compareTo(CCBitCounts other) {
	if (this.bitCount.equals(other.bitCount)) {
	    return this.contextChar.compareTo(other.contextChar);
	}
	else return this.bitCount.compareTo(other.bitCount);
    }

    public String toString() {
	return "(" + contextChar + "," + bitCount + ")";
    }
}
 
